import tushare as ts
import pymysql
import pandas as pd
import time
import sys

# 设置 Tushare 的 token
ts.set_token('d8678ccb5aa9d9922918505e441be02a1a3cf784d47ba5660cfe8016')
pro = ts.pro_api()

# 连接到 MySQL 数据库，指定字符集为 utf8mb4
conn = pymysql.connect(
    host='localhost',
    user='root',
    password='123456',
    database='test06',
    charset='utf8mb4'
)
cursor = conn.cursor()


def print_progress(iteration, total, prefix='', suffix='', decimals=1, length=50, fill='█', print_end="\r"):
    """
    显示进度条
    :param iteration: 当前迭代次数
    :param total: 总迭代次数
    :param prefix: 前缀文本
    :param suffix: 后缀文本
    :param decimals: 百分比的小数位数
    :param length: 进度条长度
    :param fill: 进度条填充字符
    :param print_end: 每次打印后的结束字符
    """
    percent = ("{0:." + str(decimals) + "f}").format(100 * (iteration / float(total)))
    filled_length = int(length * iteration // total)
    bar = fill * filled_length + '-' * (length - filled_length)
    sys.stdout.write(f'\r{prefix} |{bar}| {percent}% {suffix}')
    sys.stdout.flush()
    # 迭代完成后换行
    if iteration == total:
        print()


def get_codeList_dfcf(type_param, belong, table):
    if type_param == 'all':
        type_condition = 'IS NOT NULL'
    else:
        type_condition = f"= '{type_param}'"

    sql = f"SELECT bkdm, name, type, belong, category, type_copy, belong_copy, sort, sort_all FROM {table} WHERE type {type_condition} AND belong = '{belong}' AND (mark != '-1' OR mark IS NULL)"
    cursor.execute(sql)
    result = cursor.fetchall()
    columns = [desc[0] for desc in cursor.description]
    return [dict(zip(columns, row)) for row in result]


def get_codeList_dfcf1(type_param, belong, table):
    if belong == "ETF":
        type_str = "'industry', 'concept', 'region'"
        # 添加 sort_all 列的查询
        sql = f"SELECT bkdm, name, category, belong_copy, belong, type_copy, type, sort, sort_all FROM {table} WHERE type_copy IN ({type_str}) AND belong = '{belong}' AND (mark != '-1' OR mark IS NULL)"
    else:
        if type_param == 'all':
            type_condition = 'IS NOT NULL'
        else:
            type_condition = f"= '{type_param}'"
        sql = f"SELECT bkdm, name, type, belong, category, type_copy, belong_copy, sort, sort_all FROM {table} WHERE type {type_condition} AND belong = '{belong}' AND (mark != '-1' OR mark IS NULL) AND sort = 1"

    cursor.execute(sql)
    result = cursor.fetchall()
    columns = [desc[0] for desc in cursor.description]
    return [dict(zip(columns, row)) for row in result]


def determine_market(fund_code):
    """
    根据基金代码规则判断所属市场
    """
    if isinstance(fund_code, str):
        if fund_code.startswith('5'):
            if fund_code.startswith('508'):
                return 'SH'
            elif fund_code.startswith('50'):
                return 'SH'
            else:
                return 'SH'
        elif fund_code.startswith('15'):
            return 'SZ'
        elif fund_code.startswith('16'):
            return 'SZ'
        elif fund_code.startswith('180'):
            return 'SZ'
        else:
            return 'OF'
    else:
        print(f"基金代码 {fund_code} 不是有效的字符串，跳过")
        return None


def format_date(date_str):
    if len(date_str) == 8 and date_str.isdigit():
        return f"{date_str[:4]}-{date_str[4:6]}-{date_str[6:]}"
    return date_str


def get_daily_ETF(choice):
    if choice == 1:
        bdkm_arr = get_codeList_dfcf('index', 'ETF', 'bkxx')
    elif choice == 2:
        bdkm_arr = get_codeList_dfcf1('index', 'ETF', 'bkxx')
    else:
        print("无效的选择，请输入 1 或 2。")
        return

    total = len(bdkm_arr)
    retryList = []

    # 初始化进度条
    print_progress(0, total, prefix='数据抓取进度:', suffix='完成', length=50)

    for j in range(total):
        bkdm = bdkm_arr[j]['bkdm']
        name = bdkm_arr[j]['name']
        category = bdkm_arr[j]['category']
        if bdkm_arr[j]['belong_copy'] != bdkm_arr[j]['belong']:
            belong = bdkm_arr[j]['belong_copy']
        else:
            belong = bdkm_arr[j]['belong']
        if bdkm_arr[j]['type_copy'] != bdkm_arr[j]['type']:
            type_ = bdkm_arr[j]['type_copy']
        else:
            type_ = bdkm_arr[j]['type']
        sort = bdkm_arr[j]['sort']
        sort_all = bdkm_arr[j]['sort_all']

        # 判断市场
        market = determine_market(bkdm)
        if market is None:
            # 更新进度条
            print_progress(j + 1, total, prefix='数据抓取进度:', suffix='完成', length=50)
            continue
        bkdm_with_market = f"{bkdm}.{market}"

        try:
            # 添加延迟，这里设置为 1 秒，可以根据实际情况调整
            time.sleep(1)
            # 这里用 Tushare 模拟获取基金历史净值数据，你可按需替换接口
            data = pro.fund_nav(ts_code=bkdm_with_market)
            if data.empty:
                retryList.append(j)
            else:
                for _, row in data.iterrows():
                    date = row['nav_date']
                    accum_nav = row['accum_nav']
                    unit_nav = row['unit_nav']

                    formatted_date = format_date(date)

                    if pd.notna(accum_nav):
                        zf = float(accum_nav) / 100
                    else:
                        zf = 0
                    close = float(unit_nav) if pd.notna(unit_nav) else 0
                    close2 = float(accum_nav) if pd.notna(accum_nav) else 0

                    sql = f"INSERT IGNORE INTO `bksj` (bkdm, name, time, zf, close, close2, sign, type, belong, category, sort, sort_all) VALUES ('{bkdm}', '{name}', '{formatted_date}', {zf}, {close}, {close2}, 0, '{type_}', '{belong}', '{category}', {sort}, {sort_all})"
                    cursor.execute(sql)
                    conn.commit()
        except Exception as e:
            print(f"\n抓取 {bkdm_with_market} 数据时出错: {e}")
            retryList.append(j)

        # 更新进度条
        print_progress(j + 1, total, prefix='数据抓取进度:', suffix='完成', length=50)

    # 如果存在需要重新爬取的股票代码，则再次进行爬取
    if retryList:
        print(f"\n开始重试失败的 {len(retryList)} 个请求...")
        # 初始化重试进度条
        print_progress(0, len(retryList), prefix='重试进度:', suffix='完成', length=50)

        for i, j in enumerate(retryList):
            bkdm = bdkm_arr[j]['bkdm']
            name = bdkm_arr[j]['name']
            category = bdkm_arr[j]['category']
            if bdkm_arr[j]['belong_copy'] != bdkm_arr[j]['belong']:
                belong = bdkm_arr[j]['belong_copy']
            else:
                belong = bdkm_arr[j]['belong']
            if bdkm_arr[j]['type_copy'] != bdkm_arr[j]['type']:
                type_ = bdkm_arr[j]['type_copy']
            else:
                type_ = bdkm_arr[j]['type']
            sort = bdkm_arr[j]['sort']
            sort_all = bdkm_arr[j]['sort_all']

            # 判断市场
            market = determine_market(bkdm)
            if market is None:
                # 更新进度条
                print_progress(i + 1, len(retryList), prefix='重试进度:', suffix='完成', length=50)
                continue
            bkdm_with_market = f"{bkdm}.{market}"

            try:
                # 添加延迟，这里设置为 1 秒，可以根据实际情况调整
                time.sleep(1)
                data = pro.fund_nav(ts_code=bkdm_with_market, market='E')
                if not data.empty:
                    for _, row in data.iterrows():
                        date = row['nav_date']
                        accum_nav = row['accum_nav']
                        unit_nav = row['unit_nav']

                        formatted_date = format_date(date)

                        if pd.notna(accum_nav):
                            zf = float(accum_nav) / 100
                        else:
                            zf = 0
                        close = float(unit_nav) if pd.notna(unit_nav) else 0
                        close2 = float(accum_nav) if pd.notna(accum_nav) else 0

                        sql = f"INSERT IGNORE INTO `bksj` (bkdm, name, time, zf, close, close2, sign, type, belong, category, sort, sort_all) VALUES ('{bkdm}', '{name}', '{formatted_date}', {zf}, {close}, {close2}, 0, '{type_}', '{belong}', '{category}', {sort}, {sort_all})"
                        cursor.execute(sql)
                        conn.commit()
            except Exception as e:
                print(f"\n重新抓取 {bkdm_with_market} 数据时出错: {e}")

            # 更新重试进度条
            print_progress(i + 1, len(retryList), prefix='重试进度:', suffix='完成', length=50)

    # 数据爬取完成之后对信息进行修改
    print("\n开始更新分类信息...")
    type_str = "'industry', 'concept', 'region'"
    sql = f"SELECT * FROM bkxx WHERE belong = 'ETF' AND type_copy IN ({type_str})"
    cursor.execute(sql)
    main = cursor.fetchall()
    columns = [desc[0] for desc in cursor.description]
    main = [dict(zip(columns, row)) for row in main]

    # 初始化更新进度条
    print_progress(0, len(main), prefix='更新分类进度:', suffix='完成', length=50)

    for i, row in enumerate(main):
        bkdm = row['bkdm']
        name = row['name']
        type_ = row['type_copy']
        sort = row['sort']
        sort_all = row['sort_all']
        sql = f"UPDATE `bksj` SET type = '{type_}', sort = {sort}, sort_all = {sort_all} WHERE bkdm = '{bkdm}'"
        cursor.execute(sql)
        conn.commit()

        # 更新分类进度条
        print_progress(i + 1, len(main), prefix='更新分类进度:', suffix='完成', length=50)


def updateCategory(belong):
    type_str = "'industry', 'concept', 'region'"
    if belong == "ETF":
        sql = f"SELECT * FROM bkxx WHERE belong = '{belong}' AND type_copy IN ({type_str})"
    else:
        sql = f"SELECT * FROM bkxx WHERE belong = '{belong}' AND type IN ({type_str})"

    cursor.execute(sql)
    bkxx = cursor.fetchall()
    columns = [desc[0] for desc in cursor.description]
    bkxx = [dict(zip(columns, row)) for row in bkxx]

    # 初始化进度条
    print_progress(0, len(bkxx), prefix='更新类别进度:', suffix='完成', length=50)

    for i, row in enumerate(bkxx):
        category = row['category']
        bkdm = row['bkdm']
        sql = f"UPDATE bksj SET category = '{category}' WHERE bkdm = '{bkdm}' AND time > '2021-10-26'"
        cursor.execute(sql)
        conn.commit()

        # 更新进度条
        print_progress(i + 1, len(bkxx), prefix='更新类别进度:', suffix='完成', length=50)


if __name__ == "__main__":
    choice = int(input("请输入选择（1 （全部）或 2（排序））："))
    get_daily_ETF(choice)
    updateCategory('ETF')
    cursor.close()
    conn.close()    